{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HKT-NezBJ4Nd"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DwBljPxTJ4Ng"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hIb8SDeKJ4Nh"
      },
      "source": [
        "# Pruning for on-device inference w/ XNNPACK"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RX1fje9OJ4Ni"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/model_optimization/guide/pruning/pruning_for_on_device_inference\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/model-optimization/blob/master/tensorflow_model_optimization/g3doc/guide/pruning/pruning_for_on_device_inference.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/model-optimization/blob/master/tensorflow_model_optimization/g3doc/guide/pruning/pruning_for_on_device_inference.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/model-optimization/tensorflow_model_optimization/g3doc/guide/pruning/pruning_for_on_device_inference.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8uCVQVRMJ4Nj"
      },
      "source": [
        "Welcome to the guide on Keras weights pruning for improving latency of on-device inference via [XNNPACK](https://github.com/google/XNNPACK).\n",
        "\n",
        "This guide presents the usage of the newly introduced `tfmot.sparsity.keras.PruningPolicy` API and demonstrates how it could be used for accelerating mostly convolutional models on modern CPUs using [XNNPACK Sparse inference](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/README.md#sparse-inference).\n",
        "\n",
        "The guide covers the following steps of the model creation process:\n",
        "* Build and train the dense baseline\n",
        "* Fine-tune model with pruning\n",
        "* Convert to TFLite\n",
        "* On-device benchmark\n",
        "\n",
        "The guide doesn't cover the best practices for the fine-tuning with pruning. For more detailed information on this topic, please check out our [comprehensive guide](https://www.tensorflow.org/model_optimization/guide/pruning/comprehensive_guide.md)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UF5sWYUZJ4Nk"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "re0qdmOAJ4Nk"
      },
      "outputs": [],
      "source": [
        "! pip install -q tf-nightly\n",
        "! pip install -q tensorflow-model-optimization==0.5.1.dev0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIn7sB8-J4Nk"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "from tensorflow import keras\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eOeRh5hvJ4Nl"
      },
      "source": [
        "## Build and train the dense model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XNxUSmleJ4Nl"
      },
      "source": [
        "We build and train a simple baseline CNN for classification task on [CIFAR10](https://www.cs.toronto.edu/~kriz/cifar.html) dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ws4cmZCJJ4Nm"
      },
      "outputs": [],
      "source": [
        "# Load CIFAR10 dataset.\n",
        "(ds_train, ds_val, ds_test), ds_info = tfds.load(\n",
        "    'cifar10',\n",
        "    split=['train[:90%]', 'train[90%:]', 'test'],\n",
        "    as_supervised=True,\n",
        "    with_info=True,\n",
        ")\n",
        "\n",
        "# Normalize the input image so that each pixel value is between 0 and 1.\n",
        "def normalize_img(image, label):\n",
        "  \"\"\"Normalizes images: `uint8` -> `float32`.\"\"\"\n",
        "  return tf.image.convert_image_dtype(image, tf.float32), label\n",
        "\n",
        "# Load the data in batches of 128 images.\n",
        "batch_size = 128\n",
        "def prepare_dataset(ds, buffer_size=None):\n",
        "  ds = ds.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "  ds = ds.cache()\n",
        "  if buffer_size:\n",
        "    ds = ds.shuffle(buffer_size)\n",
        "  ds = ds.batch(batch_size)\n",
        "  ds = ds.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "  return ds\n",
        "\n",
        "ds_train = prepare_dataset(ds_train,\n",
        "                           buffer_size=ds_info.splits['train'].num_examples)\n",
        "ds_val = prepare_dataset(ds_val)\n",
        "ds_test = prepare_dataset(ds_test)\n",
        "\n",
        "# Build the dense baseline model.\n",
        "dense_model = keras.Sequential([\n",
        "    keras.layers.InputLayer(input_shape=(32, 32, 3)),\n",
        "    keras.layers.ZeroPadding2D(padding=1),\n",
        "    keras.layers.Conv2D(\n",
        "        filters=8,\n",
        "        kernel_size=(3, 3),\n",
        "        strides=(2, 2),\n",
        "        padding='valid'),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.DepthwiseConv2D(kernel_size=(3, 3), padding='same'),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.Conv2D(filters=16, kernel_size=(1, 1)),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.ZeroPadding2D(padding=1),\n",
        "    keras.layers.DepthwiseConv2D(\n",
        "        kernel_size=(3, 3), strides=(2, 2), padding='valid'),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.Conv2D(filters=32, kernel_size=(1, 1)),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.GlobalAveragePooling2D(),\n",
        "    keras.layers.Flatten(),\n",
        "    keras.layers.Dense(10)\n",
        "])\n",
        "\n",
        "# Compile and train the dense model for 10 epochs.\n",
        "dense_model.compile(\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    optimizer='adam',\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "dense_model.fit(\n",
        "  ds_train,\n",
        "  epochs=10,\n",
        "  validation_data=ds_val)\n",
        "\n",
        "# Evaluate the dense model.\n",
        "_, dense_model_accuracy = dense_model.evaluate(ds_test, verbose=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "32h7h4YYJ4Nn"
      },
      "source": [
        "## Build the sparse model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nXAXjJcuJ4Nn"
      },
      "source": [
        "Using the instructions from the [comprehensive guide](https://www.tensorflow.org/model_optimization/guide/pruning/comprehensive_guide.md), we apply `tfmot.sparsity.keras.prune_low_magnitude` function with parameters that target on-device acceleration via pruning i.e. `tfmot.sparsity.keras.PruneForLatencyOnXNNPack` policy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N1WQt5dmJ4Nn"
      },
      "outputs": [],
      "source": [
        "prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude\n",
        "\n",
        "# Compute end step to finish pruning after after 5 epochs.\n",
        "end_epoch = 5\n",
        "\n",
        "num_iterations_per_epoch = len(ds_train)\n",
        "end_step =  num_iterations_per_epoch * end_epoch\n",
        "\n",
        "# Define parameters for pruning.\n",
        "pruning_params = {\n",
        "      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.25,\n",
        "                                                               final_sparsity=0.75,\n",
        "                                                               begin_step=0,\n",
        "                                                               end_step=end_step),\n",
        "      'pruning_policy': tfmot.sparsity.keras.PruneForLatencyOnXNNPack()\n",
        "}\n",
        "\n",
        "# Try to apply pruning wrapper with pruning policy parameter.\n",
        "try:\n",
        "  model_for_pruning = prune_low_magnitude(dense_model, **pruning_params)\n",
        "except ValueError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wzDggTmYJ4No"
      },
      "source": [
        "The call `prune_low_magnitude` results in `ValueError` with the message `Could not find a GlobalAveragePooling2D layer with keepdims = True in all output branches`. The message indicates that the model isn't supported for pruning with policy `tfmot.sparsity.keras.PruneForLatencyOnXNNPack` and specifically the layer `GlobalAveragePooling2D` requires the parameter `keepdims = True`. Let's fix that and reapply `prune_low_magnitude` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qvALAbZeJ4No"
      },
      "outputs": [],
      "source": [
        "fixed_dense_model = keras.Sequential([\n",
        "    keras.layers.InputLayer(input_shape=(32, 32, 3)),\n",
        "    keras.layers.ZeroPadding2D(padding=1),\n",
        "    keras.layers.Conv2D(\n",
        "        filters=8,\n",
        "        kernel_size=(3, 3),\n",
        "        strides=(2, 2),\n",
        "        padding='valid'),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.DepthwiseConv2D(kernel_size=(3, 3), padding='same'),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.Conv2D(filters=16, kernel_size=(1, 1)),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.ZeroPadding2D(padding=1),\n",
        "    keras.layers.DepthwiseConv2D(\n",
        "        kernel_size=(3, 3), strides=(2, 2), padding='valid'),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.Conv2D(filters=32, kernel_size=(1, 1)),\n",
        "    keras.layers.BatchNormalization(),\n",
        "    keras.layers.ReLU(),\n",
        "    keras.layers.GlobalAveragePooling2D(keepdims=True),\n",
        "    keras.layers.Flatten(),\n",
        "    keras.layers.Dense(10)\n",
        "])\n",
        "\n",
        "# Use the pretrained model for pruning instead of training from scratch.\n",
        "fixed_dense_model.set_weights(dense_model.get_weights())\n",
        "\n",
        "# Try to reapply pruning wrapper.\n",
        "model_for_pruning = prune_low_magnitude(fixed_dense_model, **pruning_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51K8A0XCJ4Np"
      },
      "source": [
        "Invocation of `prune_low_magnitude` has finished without any errors meaning that the model is fully supported for the `tfmot.sparsity.keras.PruneForLatencyOnXNNPack` policy and can be accelerated using [XNNPACK Sparse inference](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/README.md#sparse-inference)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NG4yrdNUJ4Np"
      },
      "source": [
        "### Fine-tune the sparse model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EqhqM7AtJ4Np"
      },
      "source": [
        "Following the [pruning example](https://www.tensorflow.org/model_optimization/guide/pruning/pruning_with_keras.md), we fine-tune the sparse model using the weights of the dense model. We start fine-tuning of the model with 25% sparsity (25% of the weights are set to zero) and end with 75% sparsity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GzdS6AgRJ4Np"
      },
      "outputs": [],
      "source": [
        "logdir = tempfile.mkdtemp()\n",
        "\n",
        "callbacks = [\n",
        "  tfmot.sparsity.keras.UpdatePruningStep(),\n",
        "  tfmot.sparsity.keras.PruningSummaries(log_dir=logdir),\n",
        "]\n",
        "\n",
        "model_for_pruning.compile(\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    optimizer='adam',\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "model_for_pruning.fit(\n",
        "  ds_train,\n",
        "  epochs=15,\n",
        "  validation_data=ds_val,\n",
        "  callbacks=callbacks)\n",
        "\n",
        "# Evaluate the dense model.\n",
        "_, pruned_model_accuracy = model_for_pruning.evaluate(ds_test, verbose=0)\n",
        "\n",
        "print('Dense model test accuracy:', dense_model_accuracy)\n",
        "print('Pruned model test accuracy:', pruned_model_accuracy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeI9QSIMJ4Nq"
      },
      "source": [
        "The logs show the progression of sparsity on a per-layer basis."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fGDkSxKJJ4Nq"
      },
      "outputs": [],
      "source": [
        "#docs_infra: no_execute\n",
        "%tensorboard --logdir={logdir}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7hJ5PX_y7psT"
      },
      "source": [
        "After the fine-tuning with pruning, test accuracy demonstrates a modest improvement (43% to 44%) compared to the dense model. Let's compare on-device latency using [TFLite benchmark](https://www.tensorflow.org/lite/performance/measurement)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xCRpxczJ4Nq"
      },
      "source": [
        "## Model conversion and benchmarking\n",
        "\n",
        "To convert the pruned model into TFLite, we need replace the `PruneLowMagnitude` wrappers with original layers via the `strip_pruning` function. Also, since the weights of the pruned model (`model_for_pruning`) are mostly zeros, we may apply an optimization `tf.lite.Optimize.EXPERIMENTAL_SPARSITY` to efficiently store the resulted TFLite model. This optimization flag is not required for the dense model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AAJr2XKCJ4Nr"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(dense_model)\n",
        "dense_tflite_model = converter.convert()\n",
        "\n",
        "_, dense_tflite_file = tempfile.mkstemp('.tflite')\n",
        "with open(dense_tflite_file, 'wb') as f:\n",
        "  f.write(dense_tflite_model)\n",
        "\n",
        "model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(model_for_export)\n",
        "converter.optimizations = [tf.lite.Optimize.EXPERIMENTAL_SPARSITY]\n",
        "pruned_tflite_model = converter.convert()\n",
        "\n",
        "_, pruned_tflite_file = tempfile.mkstemp('.tflite')\n",
        "with open(pruned_tflite_file, 'wb') as f:\n",
        "  f.write(pruned_tflite_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uLClpRYq7psT"
      },
      "source": [
        "Following the instructions of [TFLite Model Benchmarking Tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/benchmark), we build the tool, upload it to the Android device together with dense and pruned TFLite models, and benchmark both models on the device."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-qamJwAM7psU"
      },
      "outputs": [],
      "source": [
        "! adb shell /data/local/tmp/benchmark_model \\\n",
        "    --graph=/data/local/tmp/dense_model.tflite \\\n",
        "    --use_xnnpack=true \\\n",
        "    --num_runs=100 \\\n",
        "    --num_threads=1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fpTxyOcd7psU"
      },
      "outputs": [],
      "source": [
        "! adb shell /data/local/tmp/benchmark_model \\\n",
        "    --graph=/data/local/tmp/pruned_model.tflite \\\n",
        "    --use_xnnpack=true \\\n",
        "    --num_runs=100 \\\n",
        "    --num_threads=1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fKOVVoYD7psU"
      },
      "source": [
        "Benchmarks on Pixel 4 resulted in average inference time of *17us* for the dense model and *12us* for the pruned model. The on-device benchmarks demonstrate a clear **5us** or **30%** improvements in latency even for such small models. In our experience, larger models based on [MobileNetV3](https://www.tensorflow.org/api_docs/python/tf/keras/applications/mobilenet_v3) or [EfficientNet-lite](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite) show similar performance improvements. The speed-up varies based on the relative contribution of 1x1 convolutions to the overall model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YqL_Nhuw7psU"
      },
      "source": [
        "## Conclusion"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "03Icj3vV7psU"
      },
      "source": [
        "In this tutorial, we show how one may create sparse models for faster on-device performance using the new functionality introduced by the TF MOT API and XNNPack. These sparse models are smaller and faster than their dense counterparts while retaining or even surpassing their quality.\n",
        "\n",
        "We encourage you to try this new capability which can be particularly important for deploying your models on device."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "pruning_for_on_device_inference.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
